library(data.table)
library(OpenImageR)
library(magrittr)
library(pROC)
library(mxnet)
library(abind)

# Read data

train_dat <- fread('train.csv', data.table = FALSE)
test_dat <- fread('test.csv', data.table = FALSE)
submit_dat <- fread('sample_submission.csv', data.table = FALSE)
label_dat <- fread('label/label_list.csv', data.table = FALSE)

valid_dat <- train_dat[train_dat[,'PID'] %in% tail(unique(train_dat[,'PID']), n = 2),]
train_dat <- train_dat[!train_dat[,'PID'] %in% valid_dat[,'PID'],]

# Settings

fold_pos <- 50L

img_dir <- 'image/'
process_dir <- 'processed/'
label_dir <- 'label/'
PIDs <- unique(train_dat[,'PID'])

neg_pixel <- c(461, 465, 238, 241, 51, 54, 29, 89, 133)
neg_pixel <- neg_pixel * fold_pos

## Iterator

my_iterator_core <- function (batch_size, random_crop = 224) {
  
  batch <- 0
  batch_per_epoch <- nrow(train_dat) / batch_size
  
  reset <- function() {batch <<- 0}
  
  iter.next <- function() {
    batch <<- batch+1
    if (batch > batch_per_epoch) {return(FALSE)} else {return(TRUE)}
  }
  
  value <- function() {
    batch_PID <- sample(PIDs, 1, replace = TRUE)
    load(paste0(process_dir, batch_PID, '.RData'))
    batch_idxID <- sample(1:dim(x_array)[3], batch_size, replace = TRUE)
    batch_x_list <- list()
    batch_y_list <- list()
    batch_mask_list <- list()
    for (j in 1:batch_size) {
      random_row <- sample(0:(dim(x_array)[1] - random_crop), 1)
      random_col <- sample(0:(dim(x_array)[2] - random_crop), 1)
      batch_x_list[[j]] <- x_array[random_row+1:random_crop,random_col+1:random_crop,batch_idxID[j],]
      batch_y_list[[j]] <- y_array[random_row+1:random_crop,random_col+1:random_crop,batch_idxID[j],]
      sample_mask_list <- list()
      for (k in 1:length(neg_pixel)) {
        sub_mask_array <- array(0L, dim = c(random_crop, random_crop))
        neg_mask_pos <- sample(1:length(sub_mask_array), neg_pixel[k])
        sub_mask_array[neg_mask_pos] <- 1L
        pos_mask_pos <- which(y_array[random_row+1:random_crop,random_col+1:random_crop,batch_idxID[j],k] == 1)
        sub_mask_array[pos_mask_pos] <- fold_pos
        mirror_mask_pos <- which(y_array[random_row+1:random_crop,(dim(y_array)[2] + 1) - (random_col+1:random_crop),batch_idxID[j],k] == 1)
        sub_mask_array[mirror_mask_pos] <- fold_pos
        sample_mask_list[[k]] <- sub_mask_array
      }
      batch_mask_list[[j]] <- abind(sample_mask_list, along = 3)
    }
    batch_x <- abind(batch_x_list, along = 4)
    batch_y <- abind(batch_y_list, along = 4)
    batch_mask <- abind(batch_mask_list, along = 4)
    data <- mx.nd.array(batch_x)
    label <- mx.nd.array(batch_y)
    mask <- mx.nd.array(batch_mask)
    return(list(data = data, label = label, mask = mask))
  }
  
  return(list(reset = reset, iter.next = iter.next, value = value, batch_size = batch_size, batch = batch))
}

my_iterator_func <- setRefClass("Custom_Iter",
                                fields = c("iter", "batch_size", "random_crop"),
                                contains = "Rcpp_MXArrayDataIter",
                                methods = list(
                                  initialize = function (iter, batch_size = 8, random_crop = 224){
                                    .self$iter <- my_iterator_core(batch_size = batch_size, random_crop = random_crop)
                                    .self
                                  },
                                  value = function () {
                                    .self$iter$value()
                                  },
                                  iter.next = function () {
                                    .self$iter$iter.next()
                                  },
                                  reset = function () {
                                    .self$iter$reset()
                                  },
                                  finalize = function () {
                                  }
                                )
)

my_iter <- my_iterator_func(iter = NULL, batch_size = 16, random_crop = 224)

## Model architecture

data <- mx.symbol.Variable('data')
bn_data <- mx.symbol.BatchNorm(data = data, fix.gamma = TRUE, name = 'bn_data')

down_conv1 <- mx.symbol.Convolution(data = bn_data, kernel = c(7, 7), pad = c(3, 3), num_filter = 32, no.bias = TRUE, name = 'down_conv1')
down_convbn1 <- mx.symbol.BatchNorm(data = down_conv1, fix.gamma = FALSE, name = 'down_convbn1')
down_convrelu1 <- mx.symbol.Activation(data = down_convbn1, act_type = "relu", name = 'down_convrelu1')
down_convpool1 <- mx.symbol.Pooling(data = down_convrelu1, pool_type = "avg", kernel = c(2, 2), stride = c(2, 2), name = 'down_convpool1')

down_conv2 <- mx.symbol.Convolution(data = down_convpool1, kernel = c(5, 5), pad = c(2, 2), num_filter = 64, no.bias = TRUE, name = 'down_conv2')
down_convbn2 <- mx.symbol.BatchNorm(data = down_conv2, fix.gamma = FALSE, name = 'down_convbn2')
down_convrelu2 <- mx.symbol.Activation(data = down_convbn2, act_type = "relu", name = 'down_convrelu2')
down_convpool2 <- mx.symbol.Pooling(data = down_convrelu2, pool_type = "avg", kernel = c(2, 2), stride = c(2, 2), name = 'down_convpool2')

down_conv3 <- mx.symbol.Convolution(data = down_convpool2, kernel = c(3, 3), pad = c(1, 1), num_filter = 128, no.bias = TRUE, name = 'down_conv3')
down_convbn3 <- mx.symbol.BatchNorm(data = down_conv3, fix.gamma = FALSE, name = 'down_convbn3')
down_convrelu3 <- mx.symbol.Activation(data = down_convbn3, act_type = "relu", name = 'down_convrelu3')
down_convpool3 <- mx.symbol.Pooling(data = down_convrelu3, pool_type = "avg", kernel = c(2, 2), stride = c(2, 2), name = 'down_convpool3')

down_conv4 <- mx.symbol.Convolution(data = down_convpool3, kernel = c(3, 3), pad = c(1, 1), num_filter = 256, no.bias = TRUE, name = 'down_conv4')
down_convbn4 <- mx.symbol.BatchNorm(data = down_conv4, fix.gamma = FALSE, name = 'down_convbn4')
down_convrelu4 <- mx.symbol.Activation(data = down_convbn4, act_type = "relu", name = 'down_convrelu4')
down_convpool4 <- mx.symbol.Pooling(data = down_convrelu4, pool_type = "avg", kernel = c(2, 2), stride = c(2, 2), name = 'down_convpool4')

down_conv5 <- mx.symbol.Convolution(data = down_convpool4, kernel = c(3, 3), pad = c(1, 1), num_filter = 512, no.bias = TRUE, name = 'down_conv5')
down_convbn5 <- mx.symbol.BatchNorm(data = down_conv5, fix.gamma = FALSE, name = 'down_convbn5')
down_convrelu5 <- mx.symbol.Activation(data = down_convbn5, act_type = "relu", name = 'down_convrelu5')
down_convpool5 <- mx.symbol.Pooling(data = down_convrelu5, pool_type = "avg", kernel = c(2, 2), stride = c(2, 2), name = 'down_convpool5')

down_conv6 <- mx.symbol.Convolution(data = down_convpool5, kernel = c(3, 3), pad = c(1, 1), num_filter = 512, no.bias = TRUE, name = 'down_conv6')
down_convbn6 <- mx.symbol.BatchNorm(data = down_conv6, fix.gamma = FALSE, name = 'down_convbn6')
down_convrelu6 <- mx.symbol.Activation(data = down_convbn6, act_type = "relu", name = 'down_convrelu6')

up_deconv1 <- mx.symbol.Deconvolution(data = down_convrelu6, kernel = c(2, 2), stride = c(2, 2), num_filter = 512, name = 'up_deconv1')
up_deconvbn1 <- mx.symbol.BatchNorm(data = up_deconv1, fix.gamma = FALSE, name = 'up_deconvbn1')
up_deconvrelu1 <- mx.symbol.Activation(data = up_deconvbn1, act_type = "relu", name = 'up_deconvrelu1')
up_conv1 <- mx.symbol.Convolution(data = up_deconvrelu1, kernel = c(3, 3), pad = c(1, 1), num_filter = 512, no.bias = TRUE, name = 'up_conv1')
up_convbn1 <- mx.symbol.BatchNorm(data = up_conv1, fix.gamma = FALSE, name = 'up_convbn1')
up_convrelu1 <- mx.symbol.Activation(data = up_convbn1, act_type = "relu", name = 'up_convrelu1')

up_concat2 <- mx.symbol.concat(data = list(up_convrelu1, down_convrelu5), num.args = 2, dim = 1, name = 'up_concat2')
up_deconv2 <- mx.symbol.Deconvolution(data = up_concat2, kernel = c(2, 2), stride = c(2, 2), num_filter = 256, name = 'up_deconv2')
up_deconvbn2 <- mx.symbol.BatchNorm(data = up_deconv2, fix.gamma = FALSE, name = 'up_deconvbn2')
up_deconvrelu2 <- mx.symbol.Activation(data = up_deconvbn2, act_type = "relu", name = 'up_deconvrelu2')
up_conv2 <- mx.symbol.Convolution(data = up_deconvrelu2, kernel = c(3, 3), pad = c(1, 1), num_filter = 256, no.bias = TRUE, name = 'up_conv2')
up_convbn2 <- mx.symbol.BatchNorm(data = up_conv2, fix.gamma = FALSE, name = 'up_convbn2')
up_convrelu2 <- mx.symbol.Activation(data = up_convbn2, act_type = "relu", name = 'up_convrelu2')

up_concat3 <- mx.symbol.concat(data = list(up_convrelu2, down_convrelu4), num.args = 2, dim = 1, name = 'up_concat3')
up_deconv3 <- mx.symbol.Deconvolution(data = up_concat3, kernel = c(2, 2), stride = c(2, 2), num_filter = 256, name = 'up_deconv3')
up_deconvbn3 <- mx.symbol.BatchNorm(data = up_deconv3, fix.gamma = FALSE, name = 'up_deconvbn3')
up_deconvrelu3 <- mx.symbol.Activation(data = up_deconvbn3, act_type = "relu", name = 'up_deconvrelu3')
up_conv3 <- mx.symbol.Convolution(data = up_deconvrelu3, kernel = c(3, 3), pad = c(1, 1), num_filter = 256, no.bias = TRUE, name = 'up_conv3')
up_convbn3 <- mx.symbol.BatchNorm(data = up_conv3, fix.gamma = FALSE, name = 'up_convbn3')
up_convrelu3 <- mx.symbol.Activation(data = up_convbn3, act_type = "relu", name = 'up_convrelu3')

up_concat4 <- mx.symbol.concat(data = list(up_convrelu3, down_convrelu3), num.args = 2, dim = 1, name = 'up_concat4')
up_deconv4 <- mx.symbol.Deconvolution(data = up_concat4, kernel = c(2, 2), stride = c(2, 2), num_filter = 256, name = 'up_deconv4')
up_deconvbn4 <- mx.symbol.BatchNorm(data = up_deconv4, fix.gamma = FALSE, name = 'up_deconvbn4')
up_deconvrelu4 <- mx.symbol.Activation(data = up_deconvbn4, act_type = "relu", name = 'up_deconvrelu4')
up_conv4 <- mx.symbol.Convolution(data = up_deconvrelu4, kernel = c(3, 3), pad = c(1, 1), num_filter = 288, no.bias = TRUE, name = 'up_conv4')
up_convbn4 <- mx.symbol.BatchNorm(data = up_conv4, fix.gamma = FALSE, name = 'up_convbn4')
up_convrelu4 <- mx.symbol.Activation(data = up_convbn4, act_type = "relu", name = 'up_convrelu4')

up_deconv5 <- mx.symbol.Deconvolution(data = up_convrelu4, kernel = c(2, 2), stride = c(2, 2), num_filter = 288, num_group = 9, name = 'up_deconv5')
up_deconvbn5 <- mx.symbol.BatchNorm(data = up_deconv5, fix.gamma = FALSE, name = 'up_deconvbn5')
up_deconvrelu5 <- mx.symbol.Activation(data = up_deconvbn5, act_type = "relu", name = 'up_deconvrelu5')
up_conv5 <- mx.symbol.Convolution(data = up_deconvrelu5, kernel = c(3, 3), pad = c(1, 1), num_filter = 288, num_group = 9, no.bias = TRUE, name = 'up_conv5')
up_convbn5 <- mx.symbol.BatchNorm(data = up_conv5, fix.gamma = FALSE, name = 'up_convbn5')
up_convrelu5 <- mx.symbol.Activation(data = up_convbn5, act_type = "relu", name = 'up_convrelu5')

linear_pred <- mx.symbol.Convolution(data = up_convrelu5, kernel = c(3, 3), pad = c(1, 1), num_filter = 9, name = 'linear_pred')
logistic_pred <- mx.symbol.Activation(data = linear_pred, act.type = 'sigmoid', name = 'logistic_pred')

## CE loss

label <- mx.symbol.Variable(name = 'label')
mask <- mx.symbol.Variable(name = 'mask')

eps <- 1e-8
ce_loss_pos <- mx.symbol.broadcast_mul(mx.symbol.log(logistic_pred + eps), label)
ce_loss_neg <- mx.symbol.broadcast_mul(mx.symbol.log(1 - logistic_pred + eps), 1 - label)
ce_loss_sum <- mx.symbol.broadcast_add(ce_loss_pos, ce_loss_neg)
ce_loss_weight <- mx.symbol.broadcast_mul(ce_loss_sum, mask)
sum_mask <- mx.symbol.sum(mask, axis = 0:3)
sum_ce_loss <- mx.symbol.sum(ce_loss_weight, axis = 0:3)
mean_ce_loss <- mx.symbol.broadcast_div(sum_ce_loss, sum_mask)
ce_loss <- mx.symbol.MakeLoss(0 - mean_ce_loss, name = 'ce_loss')

# Training

my_optimizer <- mx.opt.create(name = "sgd", learning.rate = 0.005, momentum = 0.9, wd = 1e-4)

my.eval.metric.loss <- mx.metric.custom(
  name = "ce-loss", 
  function(real, pred) {
    return(as.array(pred))
  }
)

mx.set.seed(0)

model <- mx.model.FeedForward.create(symbol = ce_loss, X = my_iter, optimizer = my_optimizer,
                                     eval.metric = my.eval.metric.loss,
                                     input.names = c('data'), output.names = c('label', 'mask'),
                                     array.batch.size = 16, ctx = mx.gpu(), num.round = 100)

# Save and load model

model$symbol <- logistic_pred
mx.model.save(model, 'baseline (optimal)', 0)

model <- mx.model.load('baseline (optimal)', 0)

# Get optimal cuts

## Get valid predictions

valid_x_array <- array(0, dim = c(256, 256, 3, nrow(valid_dat)))

for (i in 1:nrow(valid_dat)) {
  img_path <- paste0(img_dir, valid_dat[i, 'PID'], '/', valid_dat[i, 'UID'], '.png')
  valid_x_array[,,,i] <- readImage(img_path)
}

valid_pred_array <- predict(model, valid_x_array, ctx = mx.cpu())

## Get each annotation

each_cuts <- NULL

for (j in 1:nrow(label_dat)) {
  
  ## Get y info
  
  valid_y_array <- array(0, dim = c(256, 256, 1, nrow(valid_dat)))
  
  for (i in 1:nrow(valid_dat)) {
    
    current_dir <- paste0(label_dir, valid_dat[i, 'PID'], '/', valid_dat[i, 'UID'], '/')
    anno_path <- paste0(current_dir, label_dat[j,'structure'], '.tiff')
    
    if (file.exists(anno_path)) {
      
      anno_img <- readImage(anno_path)
      anno_img <- (anno_img > 0.5) + 0L
      valid_y_array[,,,i] <- anno_img
      
    }
    
  }
  
  ## PRROC curve analysis
  
  pred_vec <- as.numeric(valid_pred_array[,,j,])
  y_vec <- as.numeric(valid_y_array)
  prev <- mean(y_vec)
  
  roc_curve <- roc(y_vec ~ pred_vec)
  roc_curve$ppv <- roc_curve$sensitivities * prev / (roc_curve$sensitivities * prev + (1 - roc_curve$specificities) * (1 - prev))
  roc_curve$f1 <- 2 * roc_curve$ppv * roc_curve$sensitivities / (roc_curve$ppv + roc_curve$sensitivities + 1e-8)
  
  best_pos <- which.max(roc_curve$f1)
  each_cuts[j] <- roc_curve$thresholds[best_pos]
  
  message('Best cut [', label_dat[j,'structure'], ']: ', formatC(each_cuts[j], format = 'f', 4))
  
}

# Get predictions

# Generate image array

test_x_array <- array(0, dim = c(256, 256, 3, nrow(test_dat)))

for (i in 1:nrow(test_dat)) {
  img_path <- paste0(img_dir, test_dat[i, 'PID'], '/', test_dat[i, 'UID'], '.png')
  test_x_array[,,,i] <- readImage(img_path)
}

# Predicting

test_y_array <- predict(model, test_x_array, ctx = mx.cpu())

# Processing

submit_dat_list <- list()

for (i in 1:nrow(test_dat)) {
  
  for (j in 1:nrow(label_dat)) {
    
    cut_p <- each_cuts[j]
    
    sub_y_array <- test_y_array[,,j,i]
    pos_info <- which(sub_y_array > cut_p, arr.ind = TRUE)
    
    if (nrow(pos_info) > 0) {
      
      sub_submit_dat <- data.frame(PID = test_dat[i,'PID'],
                                   UID = test_dat[i,'UID'],
                                   row = pos_info[,'row'],
                                   col = pos_info[,'col'],
                                   label = label_dat[j,'structure'],
                                   stringsAsFactors = FALSE)
      
      submit_dat_list[[length(submit_dat_list) + 1]] <- sub_submit_dat
      
    }
    
  }
  
}

my_submit_dat <- do.call('rbind', submit_dat_list)
write.csv(my_submit_dat, file = 'my_submission.csv', na = '', row.names = FALSE, quote = FALSE)